from langchain.agents import create_react_agent, AgentExecutor
from langchain_community.llms.tongyi import Tongyi
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import Tool
from rest_framework.response import Response
from rest_framework.views import APIView

llm = Tongyi()


def tool1(input):
    return "咨询问题"


def tool2(input):
    return "订单处理"


tools = [
    Tool(func=tool1, name="咨询问题", description="咨询问题"),
    Tool(func=tool2, name="订单处理", description="订单处理")
]

template = '''请尽可能回答以下问题。您可以使用以下工具:

           {tools}

           Use the following format:

           Question: the input question you must answer
           Thought: you should always think about what to do
           Action: the action to take, should be one of [{tool_names}]
           Action Input: the input to the action
           Observation: the result of the action
           ... (this Thought/Action/Action Input/Observation can repeat N times)
           Thought: I now know the final answer
           Final Answer: the final answer to the original input question

           Begin!

           Question: {input}
           Thought:{agent_scratchpad}'''

prompt = ChatPromptTemplate.from_template(template)
agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)


class Chain(APIView):
    def post(self, request):
        input_text = request.data.get('input')
        res = agent_executor.invoke({"input": input_text})
        return Response({"msg": 'ok', "data": res})
